from langgraph.graph import StateGraph,START,END
from typing import TypedDict

class InputState(TypedDict):
    question:str

class OutputState(TypedDict):
    answer:str

class OverallState(InputState,OutputState):
    pass


def agent_node(state:InputState):
    print("我是一个AI agent")
    return {"question":state["question"]}

def action_node(state:InputState):
    print("我现在是一个执行者。")
    step = state["question"]
    return {"answer":f"我接收到的问题是{step},读取成功了"}

# 明确指定它的输入和输出数据的结构或模式
builder = StateGraph(state_schema=OverallState,input=InputState,output=OutputState)

# 添加点
builder.add_node("agent_node",agent_node)
builder.add_node("action_node",action_node)

# 添加边
builder.add_edge(START,"agent_node")
builder.add_edge("agent_node","action_node")
builder.add_edge("action_node",END)

# 编译图
graph=builder.compile()

res = graph.invoke({"question":"今天天气怎么样"})
print(res)
